import uuid
import json
import urllib.request
import urllib.parse
import asyncio
import os
import random
import aiohttp

server_address = os.getenv("COMFYUI_SERVER_ADDRESS", "http://36.111.149.94:8188")
client_id = str(uuid.uuid4())


# 异步提交 prompt
async def queue_prompt_async(prompt):
    p = {"prompt": prompt, "client_id": client_id}
    data = json.dumps(p).encode("utf-8")
    
    async with aiohttp.ClientSession() as session:
        async with session.post(f"{server_address}/prompt", data=data) as response:
            return await response.json()
        
# docker run -itd -p 18000:8000 -e COMFYUI_SERVER_ADDRESS=http://10.0.0.4:8188  comfy-ui-api:20250708

# 异步轮询 /history 查看是否完成
async def get_history_async(prompt_id):
    url = f"{server_address}/history/{prompt_id}"
    async with aiohttp.ClientSession() as session:
        async with session.get(url) as response:
            return await response.json()


# 异步下载图片
async def get_image_async(filename, subfolder, folder_type):
    data = {"filename": filename, "subfolder": subfolder, "type": folder_type}
    url_values = urllib.parse.urlencode(data)
    
    async with aiohttp.ClientSession() as session:
        async with session.get(f"{server_address}/view?{url_values}") as response:
            return await response.read()


async def generate_image_async(positive_prompt="一只可爱的小猫在花园里玩耍", 
                               negative_prompt="", 
                               width=1024, 
                               height=1024, 
                               steps=16, 
                               cfg=1, 
                               seed=None,
                               filename_prefix="ComfyUI"):
    """
    异步生成图像的函数
    
    Args:
        positive_prompt (str): 正面提示词
        negative_prompt (str): 负面提示词
        width (int): 图像宽度
        height (int): 图像高度
        steps (int): 采样步数
        cfg (float): CFG值
        seed (int): 随机种子，如果为None则自动生成
        filename_prefix (str): 文件名前缀
    
    Returns:
        list: 生成的图像文件名列表
    """
    # 如果没有提供seed，则生成随机seed
    if seed is None:
        seed = random.randint(1, 1000000000000000)
    
    # 生成随机ID用于文件名
    random_id = str(uuid.uuid4())[:8]  # 使用UUID的前8位
    unique_filename_prefix = f"{filename_prefix}_{random_id}"
    
    # 构建 prompt
    prompt = {
        "3": {
            "inputs": {
                "seed": seed,
                "steps": steps,
                "cfg": cfg,
                "sampler_name": "lcm",
                "scheduler": "normal",
                "denoise": 1,
                "model": ["70", 0],
                "positive": ["16", 0],
                "negative": ["40", 0],
                "latent_image": ["53", 0],
            },
            "class_type": "KSampler",
            "_meta": {"title": "K采样器"},
        },
        "8": {
            "inputs": {"samples": ["3", 0], "vae": ["55", 0]},
            "class_type": "VAEDecode",
            "_meta": {"title": "VAE解码"},
        },
        "9": {
            "inputs": {"filename_prefix": unique_filename_prefix, "images": ["8", 0]},
            "class_type": "SaveImage",
            "_meta": {"title": "保存图像"},
        },
        "16": {
            "inputs": {"text": positive_prompt, "clip": ["54", 0]},
            "class_type": "CLIPTextEncode",
            "_meta": {"title": "Positive Prompt"},
        },
        "40": {
            "inputs": {"text": negative_prompt, "clip": ["54", 0]},
            "class_type": "CLIPTextEncode",
            "_meta": {"title": "Negative Prompt"},
        },
        "53": {
            "inputs": {"width": width, "height": height, "batch_size": 1},
            "class_type": "EmptySD3LatentImage",
            "_meta": {"title": "空Latent图像（SD3）"},
        },
        "54": {
            "inputs": {
                "clip_name1": "clip_l_hidream.safetensors",
                "clip_name2": "clip_g_hidream.safetensors",
                "clip_name3": "t5xxl_fp8_e4m3fn_scaled.safetensors",
                "clip_name4": "llama_3.1_8b_instruct_fp8_scaled.safetensors",
            },
            "class_type": "QuadrupleCLIPLoader",
            "_meta": {"title": "QuadrupleCLIPLoader"},
        },
        "55": {
            "inputs": {"vae_name": "ae.safetensors"},
            "class_type": "VAELoader",
            "_meta": {"title": "加载VAE"},
        },
        "70": {
            "inputs": {"shift": 3.0000000000000004, "model": ["76", 0]},
            "class_type": "ModelSamplingSD3",
            "_meta": {"title": "采样算法（SD3）"},
        },
        "76": {
            "inputs": {
                "unet_name": "hidream_i1_fast_fp8.safetensors",
                "weight_dtype": "default",
            },
            "class_type": "UNETLoader",
            "_meta": {"title": "UNet加载器"},
        },
    }

    # 提交任务
    print("Submitting prompt...")
    resp = await queue_prompt_async(prompt)
    prompt_id = resp["prompt_id"]
    print("Prompt ID:", prompt_id)

    # 异步轮询直到执行完成
    print("Waiting for result...")
    while True:
        try:
            history = await get_history_async(prompt_id)
            if prompt_id in history:
                break
            await asyncio.sleep(1)  # 使用异步sleep
        except Exception as e:
            print(f"轮询历史记录时出错: {e}")
            await asyncio.sleep(2)  # 如果出错，等待2秒后重试

    # 提取图片信息
    images = []
    outputs = history[prompt_id]["outputs"]
    for node in outputs:
        node_output = outputs[node]
        if "images" in node_output:
            for image_info in node_output["images"]:
                try:
                    img_data = await get_image_async(
                        image_info["filename"], image_info["subfolder"], image_info["type"]
                    )
                    # 生成带随机ID的本地文件名
                    original_filename = image_info["filename"]
                    file_extension = os.path.splitext(original_filename)[1]
                    local_filename = f"images/{unique_filename_prefix}_{len(images)+1}{file_extension}"
                    
                    # 确保images目录存在
                    os.makedirs("images", exist_ok=True)
                    
                    # 保存为本地文件
                    with open(local_filename, "wb") as f:
                        f.write(img_data)
                    print("Saved:", local_filename)
                    images.append(local_filename)
                except Exception as e:
                    print(f"下载图像时出错: {e}")
                    continue

    print("All done.")
    return images


# 向后兼容的同步包装函数
def generate_image(*args, **kwargs):
    """
    同步版本的生成图像函数，用于向后兼容
    """
    return asyncio.run(generate_image_async(*args, **kwargs))
